// Package middleware 防注入模块
package middleware

import (
	"bytes"
	"encoding/json"
	"io"
	"log"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

// 设置模式intercept拦截记录，observe观察记录
var wafModel = ""

// 增强版SQL注入检测正则
var sqlRegex = regexp.MustCompile(
	`(?i)(` +
		`(\b(union|select|insert|update|delete|drop|alter|truncate|exec|declare)\b\s*(\(|\s+[\w\p{L}]+))|` +
		`(\b(and|or)\b\s*(1=1|[\w\p{L}]+(=|>|<)|\s*select\b))|` +
		`(['"` + "`" + `]\s*?(--|#|/\*))|` +
		`(;\s*(drop|alter|exec))|` +
		`(/\*.*?\*/)|` +
		`(%0a|%0d|%09|\+)\s*(select|union)` +
		`)`)
var sqlLogger *log.Logger

func init() {
	// 初始化日志文件
	file, _ := os.OpenFile("tmp/AntiSQLInject.log", os.O_APPEND|os.O_CREATE|os.O_RDWR, 0644)
	sqlLogger = log.New(file, "", 0)
}

func AntiSQLInjectMiddleWare(model string) gin.HandlerFunc {
	if model != "observe" {
		model = "intercept"
	}
	wafModel = model
	return func(c *gin.Context) {
		// 统一检查方法
		checkInput := func(input string) bool {
			// 解码URL编码字符
			decoded, err := url.QueryUnescape(input)
			if err != nil {
				decoded = input
			}
			// 检查原始输入和解码后的输入
			return sqlRegex.MatchString(input) || sqlRegex.MatchString(decoded) || sqlRegex.MatchString(strings.ToLower(decoded))
		}

		// 检查URL参数
		for key, values := range c.Request.URL.Query() {
			for _, v := range values {
				if checkInput(v) {
					blockRequest(c, "URL参数", key, v)
					return
				}
			}
		}

		// 检查POST表单参数
		if err := c.Request.ParseForm(); err == nil {
			for key, values := range c.Request.PostForm {
				for _, v := range values {
					if checkInput(v) {
						blockRequest(c, "POST参数", key, v)
						return
					}
				}
			}
		}

		// 检查Cookie
		for _, cookie := range c.Request.Cookies() {
			if checkInput(cookie.Value) {
				blockRequest(c, "Cookie", cookie.Name, cookie.Value)
				return
			}
		}

		// 检查JSON请求体
		if c.Request.Header.Get("Content-Type") == "application/json" {
			body, err := io.ReadAll(c.Request.Body)
			if err == nil && len(body) > 0 {
				// 恢复请求体供后续使用
				c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
				if checkJSON(body) {
					blockRequest(c, "JSON内容", "", string(body))
					return
				}
			}
		}

		c.Next()
	}
}

// 拦截请求统一处理
func blockRequest(c *gin.Context, detectType, key, value string) {
	// 构造日志条目
	logEntry := []string{
		time.Now().Format(time.RFC3339),
		c.Request.URL.String(),
		c.Request.Method,
		c.ClientIP(),
		c.Request.Header.Get("User-Agent"),
		detectType,
		key + "=" + strings.ReplaceAll(value, "\n", "\\n"),
	}

	// 写入日志文件
	sqlLogger.Println(strings.Join(logEntry, " | "))
	if wafModel == "intercept" {
		c.AbortWithStatusJSON(403, gin.H{
			"code": 403,
			"mes":  "请求包含可疑输入特征",
		})
	}
}

// 递归检查JSON数据
func checkJSON(body []byte) bool {
	var data interface{}
	if err := json.Unmarshal(body, &data); err != nil {
		return false
	}
	return traverseJSON(data)
}

func traverseJSON(value interface{}) bool {
	switch v := value.(type) {
	case string:
		// 检查解码前后的两种形式
		decoded, _ := url.QueryUnescape(v)
		return sqlRegex.MatchString(v) || sqlRegex.MatchString(decoded) || sqlRegex.MatchString(strings.ToLower(decoded))
	case map[string]interface{}:
		for _, val := range v {
			if traverseJSON(val) {
				return true
			}
		}
	case []interface{}:
		for _, val := range v {
			if traverseJSON(val) {
				return true
			}
		}
	}
	return false
}
